"""
股票训练数据集生成模块

该模块负责从MySQL数据库读取股票数据，并生成用于模型训练的滑动窗口数据集。
主要功能包括：
1. 从数据库读取股票历史数据
2. 生成滑动窗口特征序列
3. 保存为CSV格式的训练数据集
"""
import pandas as pd
from sqlalchemy import create_engine, text
from config import DB_CONFIG
import os

def generate_training_dataset(window_size=6, output_file='./dataset/training_dataset.csv'):
    """
    生成用于模型训练的股票数据集
    
    参数:
        window_size (int, optional): 滑动窗口大小，默认6天
        output_file (str, optional): 输出CSV文件路径，默认'./dataset/training_dataset.csv'
    
    返回:
        None: 结果直接保存到CSV文件
    
    处理流程:
        1. 连接数据库并读取股票数据
        2. 按时间排序并转换日期格式
        3. 生成滑动窗口特征序列
        4. 保存为CSV格式的训练数据集
    """
    # 创建数据库引擎
    engine = create_engine(
        f"mysql+mysqlconnector://{DB_CONFIG['user']}:{DB_CONFIG['password']}@"
        f"{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
    )

    try:
        with engine.connect() as conn:
            # 读取按时间排序的股票数据
            df = pd.read_sql(
                """
                SELECT trade_date, open, high, low, close, pre_close, vol, amount 
                FROM stock_index 
                ORDER BY trade_date DESC
                """,
                conn
            )
            print(f'从数据库读取到{len(df)}条原始数据')
            if df.empty:
                print('警告：数据库查询结果为空，请检查stock_index表数据')
            
            # 转换日期格式并排序
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df = df.sort_values('trade_date').reset_index(drop=True)
            print(f'最早交易日：{df.trade_date.min()}，最晚交易日：{df.trade_date.max()}')
            
            # 生成滑动窗口数据集
            input_sequences = []  # 存储滑动窗口特征序列
            output_labels = []    # 存储对应的标签(下一天收盘价)

            
            for i in range(len(df) - window_size):
                window = df.iloc[i:i + window_size]
                next_day = df.iloc[i + window_size]
                
                # 获取窗口特征（排除结果字段）
                features = window.drop(columns=['trade_date']).values  # 获取6天的特征数据
                label = next_day['close']  # 使用第7天收盘价作为预测标签
                
                input_sequences.append(features.flatten().tolist())
                output_labels.append(label)
            
            print(f'生成有效窗口数量：{len(input_sequences)}')
            # 创建输出目录
            os.makedirs(os.path.dirname(output_file), exist_ok=True)
            print(f'生成{len(input_sequences)}条训练样本')
            print(output_labels)
            print(input_sequences)
            # 转换为DataFrame并保存
            dataset = pd.DataFrame({
                'input_features': input_sequences,
                'output_label': output_labels
            })
            
            dataset.to_csv(output_file, index=False)
            print(f'成功生成训练数据集：{output_file}')

    except Exception as e:
        print(f'数据集生成失败：{str(e)}')
        raise

if __name__ == '__main__':
    generate_training_dataset()